from functools import lru_cache
import os
from io import BytesIO
import pandas as pd
import numpy as np
import empyrical
from xqdata.constant import Frequency
from xqdata_rq import RQDataApi
from gypb.db import get_ddbapi,get_rqapi
from xq_factor import Config,add_api
from xq_factor.factor import LeafFactor, ObjectedLeafFactor, AbstractFactor,ListedFactor
from xq_factor.operator import PROPORTION, RANK, FFILL, DIFF, REF
import jinja2
from docxtpl import DocxTemplate, InlineImage
from docx.shared import Mm
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.dates as mdates

matplotlib.use("agg")
# 绘图
plt.style.use("seaborn-v0_8")

import warnings

warnings.filterwarnings("ignore")

ddbapi = get_ddbapi()
rqapi = get_rqapi()
add_api(ddbapi)
add_api(rqapi, "rq")

report_config = Config()
report_config.set_option("frequency", Frequency.DAILY)


indices = [
    "上证指数",
    "深证成指",
    "创业板指",
    "上证50",
    "沪深300",
    "中证A500",
    "中证500",
    "中证800",
    "中证1000",
    "中证2000",
    "科创50",
    "中证全指",
    "中证全债",
    "中证国债",
    "中证转债",
    "商品CFCI",
]
dict_index_code = {
    "上证指数": "000001.SSE",
    "深证成指": "399001.SZSE",
    "创业板指": "399006.SZSE",
    "上证50": "000016.SSE",
    "沪深300": "000300.SSE",
    "中证500": "000905.SSE",
    "中证800": "000906.SSE",
    "中证1000": "000852.SSE",
    "科创50": "000688.SSE",
    "中证全指": "000985.SSE",
    "中证全债": "H11001.SSE",
    "中证国债": "H11006.SSE",
    "中证转债": "000832.SSE",
    "商品CFCI": "H11061.SSE",
    "中证2000":"932000.OF",
    "国证2000":"399303.SZSE",
    "中证A500":"000510.SSE",
}


@lru_cache()
def get_indices_price(stat_date):
    indices_price = (
        rqapi.get_factor(
            "close",
            dict_index_code.values(),
            start_time="2020-01-02",
            end_time=stat_date,
        )["close"]
        .unstack()
        .rename(columns={value: key for key, value in dict_index_code.items()})
    )
    return indices_price


def get_data(prod_data, stat_date):
    report_config.set_option("start_time", "2020-01-02")
    report_config.set_option("end_time", stat_date)
    report_config.set_option("universe", prod_data.index)
    def get_secuinfo(prod_data):
        secuinfo = prod_data.copy()
        secuinfo.reset_index(inplace=True)
        secuinfo.rename(columns={"prod_code":"code","成立日期":"listed_date"},inplace=True)
        secuinfo["listed_date"] = pd.to_datetime(secuinfo["listed_date"])
        secuinfo["delisted_date"] = pd.Timestamp("2099-12-31")
        return secuinfo
    operate_date = ListedFactor(get_secuinfo(prod_data)).use_config(report_config).value
    def get_nv(name:str):
        value = LeafFactor(name).use_config(report_config).value.where(operate_date)
        # 定义处理函数：对单列只填充最后一个有效值之前的NA
        def ffill_until_last_valid(col):
            last_valid = col.last_valid_index()  # 获取该列最后一个非NA的时间索引
            if last_valid is None:  # 整列都是NA，直接返回
                return col
            mask = col.index <= last_valid  # 时间在最后一个有效值之前的位置需要填充
            return col.where(~mask, col.ffill())  # 仅在mask范围内执行ffill

        # 对DataFrame每列应用处理函数
        return value.apply(ffill_until_last_valid, axis=0)

    netvalue = get_nv("netvalue")
    cum_netvalue = get_nv("cum_netvalue")
    returns = cum_netvalue.diff() / netvalue.shift(1)
    cum_returns = (returns + 1).cumprod()

    indices_price = get_indices_price(stat_date)
    indices_returns = indices_price.pct_change()
    benchmark_returns = pd.DataFrame(
        {
            prod_code: indices_returns[benchmark]
            for prod_code, benchmark in prod_data.对标基准.items()
        },
        index=indices_returns.index,
    ).mask(returns.isna())
    benchmark_cum_returns = (benchmark_returns + 1).cumprod()
    return (
        netvalue,
        cum_netvalue,
        returns,
        cum_returns,
        benchmark_returns,
        benchmark_cum_returns,
    )


# 周度无风险收益率
rf = 0.02 / 52


def resample_index(index, freq):
    FREQ_MAP = {
        Frequency.WEEKLY: "W",
        Frequency.MONTHLY: "ME",
        Frequency.ANNUALLY: "YE",
    }
    freq_str = FREQ_MAP.get(freq, None)
    new_index = (
        pd.Series(index=index, data=index)
        .groupby(pd.Grouper(freq=freq_str))
        .last()
        .dropna()
    )
    return pd.DatetimeIndex(new_index)


def get_periods(stat_date, index: pd.DatetimeIndex, ref=1):
    return index[index <= stat_date][-1 - ref]


def trunc_returns(returns, start_time, end_time):
    mask = (returns.index > start_time) & (returns.index <= end_time)
    return returns[mask]


def get_period_returns(
    returns: pd.DataFrame,
    stat_date: pd.Timestamp,
    freq: Frequency = Frequency.WEEKLY,
    ref=1,
):
    resampled_index = resample_index(returns.index, freq)
    last_period_end = get_periods(stat_date=stat_date, index=resampled_index, ref=ref)
    return trunc_returns(returns, last_period_end, stat_date)


def calc_period_return(
    returns: pd.DataFrame,
    stat_date: pd.Timestamp,
    freq: Frequency = Frequency.WEEKLY,
    ref=1,
):
    this_period_returns = get_period_returns(
        returns=returns, stat_date=stat_date, freq=freq, ref=ref
    )
    return empyrical.cum_returns_final(this_period_returns)


def percent_format(value):
    return "{:.2%}".format(value)


# 自定义jinja2过滤器float_format，格式化输出float类型数值，保留2位小数
def float_format(value):
    return "{:.2f}".format(value)


# 自定义jinja2过滤器nv_float_format，格式化输出float类型数值，保留4位小数
def nv_float_format(value):
    return "{:.4f}".format(value)


# 自定义jinja2过滤器date2str_format，将字符串datetime"yyyy-mm-dd 00:00:00"格式化输出为"yyyy-mm-dd"字符串格式
def date2str_format(value):
    return value.strip("00:00:00")


# 自定义jinja2过滤器date2str_format，将datetime格式化输出为"yyyy-mm-dd"字符串格式
def datetime2str_format(value):
    return value.strftime("%Y-%m-%d")


def get_strategy_types(product_data):
    strategy_order = ddbapi.get_info("weekly_monthly_reports_strategy_order").set_index("strategy")["strategy_order"]
    def sort_key(s):
        # number_part = int(s.split("-")[0])  # 提取数字并转换为整数
        # return number_part
        return strategy_order.loc[s]

    # 排序
    return sorted(product_data.策略类型.unique(), key=sort_key)


def calc_product_indicators(
    netvalue,
    cum_netvalue,
    returns,
    cum_returns,
    benchmark_returns,
    benchmark_cum_returns,
    base_prod_data,
    stat_date,
):
    product_data = base_prod_data.copy()
    # this_week_returns = get_period_returns(returns, stat_date, freq=Frequency.WEEKLY)
    # this_month_returns = get_period_returns(returns, stat_date, freq=Frequency.MONTHLY)
    this_year_returns = get_period_returns(returns, stat_date, freq=Frequency.ANNUALLY)
    last_tear_returns = get_period_returns(returns, stat_date, freq=Frequency.WEEKLY, ref=52)
    # this_week_benchmark_returns = get_period_returns(benchmark_returns, stat_date, freq=Frequency.WEEKLY)
    # this_month_benchmark_returns = get_period_returns(benchmark_returns, stat_date, freq=Frequency.MONTHLY)
    this_year_benchmark_returns = get_period_returns(benchmark_returns, stat_date, freq=Frequency.ANNUALLY)
    last_tear_benchmark_returns = get_period_returns(benchmark_returns, stat_date, freq=Frequency.WEEKLY, ref=52)

    # 统一用周频数据计算波动，夏普等
    week_index = resample_index(netvalue.index, Frequency.WEEKLY)
    week_cum_returns = cum_returns.reindex(week_index)
    week_returns = week_cum_returns.pct_change()
    week_benchmark_cum_returns = benchmark_cum_returns.reindex(week_index)
    week_benchmark_returns = week_benchmark_cum_returns.pct_change()
    product_data["最新单位净值"] = netvalue.loc[stat_date]
    product_data["最新累计净值"] = cum_netvalue.loc[stat_date]

    product_data["this_week_cumulative_return"] = calc_period_return(returns, stat_date, freq=Frequency.WEEKLY)
    product_data["bm_this_week_cumulative_return"] = calc_period_return(benchmark_returns, stat_date, freq=Frequency.WEEKLY)
    product_data["this_month_cumulative_return"] = calc_period_return(returns, stat_date, freq=Frequency.MONTHLY)
    product_data["bm_this_month_cumulative_return"] = calc_period_return(benchmark_returns, stat_date, freq=Frequency.MONTHLY)    
    product_data["this_year_cumulative_return"] = calc_period_return(returns, stat_date, freq=Frequency.ANNUALLY)
    product_data["bm_this_year_cumulative_return"] = calc_period_return(benchmark_returns, stat_date, freq=Frequency.ANNUALLY)
    product_data["last_year_cumulative_return"] = calc_period_return(returns, stat_date, freq=Frequency.WEEKLY, ref=52)
    product_data["bm_last_year_cumulative_return"] = calc_period_return(benchmark_returns, stat_date, freq=Frequency.WEEKLY, ref=52)        
    product_data["total_cumulative_return"] = empyrical.cum_returns_final(returns)
    product_data["bm_total_cumulative_return"] = empyrical.cum_returns_final(benchmark_returns)

    product_data["total_maxdown"] = empyrical.max_drawdown(returns).values
    product_data["bm_total_maxdown"] = empyrical.max_drawdown(benchmark_returns).values
    product_data["this_year_maxdown"] = empyrical.max_drawdown(this_year_returns).values
    product_data["bm_this_year_maxdown"] = empyrical.max_drawdown(this_year_benchmark_returns).values
    product_data["last_year_maxdown"] = empyrical.max_drawdown(last_tear_returns).values
    product_data["bm_last_year_maxdown"] = empyrical.max_drawdown(last_tear_benchmark_returns).values    
    product_data["total_annual_return"] = empyrical.annual_return(week_returns, period="weekly")
    product_data["bm_total_annual_return"] = empyrical.annual_return(week_benchmark_returns, period="weekly")
    product_data["total_annual_alpha"] = empyrical.alpha_aligned(week_returns, week_benchmark_returns, risk_free=rf, period="weekly").values
    product_data["total_sharpe_ratio"] = empyrical.sharpe_ratio(week_returns, risk_free=rf, period="weekly")
    product_data["bm_total_sharpe_ratio"] = empyrical.sharpe_ratio(week_benchmark_returns, risk_free=rf, period="weekly")
    # product_data["total_sharpe_ratio"] = (product_data["total_annual_return"]-0.02)/(week_returns.std()*np.sqrt(52))
    # product_data["bm_total_sharpe_ratio"] = (product_data["bm_total_annual_return"]-0.02)/(week_benchmark_returns.std()*np.sqrt(52))    

    product_data["total_calmar_ratio"] = product_data["total_annual_return"] / abs(product_data["total_maxdown"])
    product_data["bm_total_calmar_ratio"] = product_data["bm_total_annual_return"] / abs(product_data["bm_total_maxdown"])

    def calc_downside_deviation(returns:pd.DataFrame):
        ret = returns.copy()
        ret[ret > 0] = 0
        downside_variance = ret.pow(2).sum() / ret.count()
        downside_deviation = np.sqrt(downside_variance)
        return downside_deviation
    product_data["total_sortino_ratio"] = empyrical.sortino_ratio(week_returns, period="weekly").values
    product_data["bm_total_sortino_ratio"] = empyrical.sortino_ratio(week_benchmark_returns, period="weekly").values
    # product_data["total_sortino_ratio"] = product_data["total_annual_return"] / (calc_downside_deviation(week_returns)*np.sqrt(52))
    # product_data["bm_total_sortino_ratio"] = product_data["bm_total_annual_return"] / (calc_downside_deviation(week_benchmark_returns)*np.sqrt(52))
    
    product_data["total_information_ratio"] = (week_returns - week_benchmark_returns).mean() / (week_returns - week_benchmark_returns).std() * np.sqrt(52)
    product_data["this_week_excess_return"] = product_data["this_week_cumulative_return"] - product_data["bm_this_week_cumulative_return"]
    product_data["need_excess"] = product_data["策略类型"].isin(["中证红利指数增强","沪深300指数增强","中证A500指数增强","中证500指数增强", "中证1000指数增强","国证2000指数增强","中证2000指数增强","量化股票多头"])
    return product_data


def calc_index_performance(stat_date):
    index_performance_data = pd.DataFrame(index=indices)
    indices_price = get_indices_price(stat_date)
    indices_returns = indices_price.pct_change()

    dict_index_type = {
        "上证指数": "A股市场",
        "深证成指": "A股市场",
        "创业板指": "A股市场",
        "上证50": "A股市场",
        "沪深300": "A股市场",
        "中证A500": "A股市场",
        "中证500": "A股市场",
        "中证800": "A股市场",
        "中证1000": "A股市场",
        "中证2000": "A股市场",
        "科创50": "A股市场",
        "中证全指": "A股市场",
        "中证全债": "债券市场",
        "中证国债": "债券市场",
        "中证转债": "债券市场",
        "商品CFCI": "大宗商品",
    }

    index_performance_data["major_asset_type"] = dict_index_type

    index_performance_data["stat_day_return"] = indices_returns.loc[stat_date]
    index_performance_data["stat_day_price"] = indices_price.loc[stat_date]
    index_performance_data["this_week_return"] = calc_period_return(
        indices_returns, stat_date, freq=Frequency.WEEKLY
    )
    index_performance_data["this_year_return"] = calc_period_return(
        indices_returns, stat_date, freq=Frequency.ANNUALLY
    )
    return index_performance_data


def get_images(product_data, cum_returns, benchmark_cum_returns, report_tpl):
    images = {}
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False

    for prod_code in cum_returns.columns:
        prod_name = product_data.loc[prod_code, "产品名称"]
        pro_bm_data = pd.DataFrame(index=cum_returns.index)
        pro_bm_data["产品走势"] = cum_returns[prod_code]
        pro_bm_data["业绩基准"] = benchmark_cum_returns[prod_code]
        pro_bm_data.plot(
            figsize=(8, 8),
            xlabel="",
            fontsize=16,
            linewidth=3.0,
            color=["firebrick", "slategrey"],
        )
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
        plt.gca().yaxis.set_major_formatter(ticker.FormatStrFormatter("%0.2f"))
        plt.grid(axis="x")
        plt.margins(x=0.015)
        plt.legend(loc="best", fontsize=16)  # loc='upper left',
        plt.title(prod_name + "净值走势图\n", fontsize=20)

        # #存储图片文件到字典
        # image_path = os.path.join(".",prod_name+'净值走势图.png')
        # plt.savefig(image_path,bbox_inches='tight') #
        # images.update({pro_code:image_path})

        # 创建一个 BytesIO 对象来保存图像数据
        image_stream = BytesIO()
        # 将 matplotlib 绘制的图像保存到 BytesIO 对象中，格式为 PNG
        plt.savefig(image_stream, format="png")
        # 将文件指针重置到文件开头，以便后续读取
        image_stream.seek(0)
        image = InlineImage(
            report_tpl, image_descriptor=image_stream, width=Mm(100), height=Mm(90)
        )
        images[prod_code] = image
        plt.close()

    return images


def generate_report(stat_date,save_dir,type="weekly"):
    base_prod_data = ddbapi.get_info("weekly_monthly_reports")
    # 生成报告
    if type == "weekly":
        template_path = "src\gypb/apps\weekly_reports/week_report.docx"
        base_prod_data = base_prod_data.query("weekly_report==True")
        file_prefix = "业绩跟踪周报"
    else:
        template_path = "src\gypb/apps\weekly_reports/month_report.docx"
        base_prod_data = base_prod_data.query("monthly_report==True")
        file_prefix = "私募月报"
    # report_tpl为报告模板
    report_tpl = DocxTemplate(template_path)

    base_prod_data.set_index("prod_code", inplace=True)
    base_prod_data.rename(
        columns={
            "prod_name": "产品名称",
            "manager_name":"私募管理人",
            "establish_date": "成立日期",
            "strategy": "策略类型",
            "sale_range": "销售范围",
            "benchmark": "对标基准",
        },
        inplace=True,
    )

    (
        netvalue,
        cum_netvalue,
        returns,
        cum_returns,
        benchmark_returns,
        benchmark_cum_returns,
    ) = get_data(base_prod_data, stat_date)

    strategy_types = get_strategy_types(base_prod_data)
    product_data = calc_product_indicators(
        netvalue,
        cum_netvalue,
        returns,
        cum_returns,
        benchmark_returns,
        benchmark_cum_returns,
        base_prod_data,
        stat_date,
    )
    index_performance_data = calc_index_performance(stat_date)
    images = get_images(base_prod_data, cum_returns, benchmark_cum_returns, report_tpl)

    context = {
        "stat_date": pd.Timestamp(stat_date),
        "strategy_types": strategy_types,
        "index_performance_data": index_performance_data,
        "images": images,
    }
    if type == "weekly":
        context["global_fund_data"] = product_data
    else:
        context["global_fund_data"] = product_data.query("销售范围=='全公司'")
        context["private_fund_data"] = product_data.query("销售范围!='全公司'")


    # 将自定义过滤器传入jinja2.Environment
    jinja_env = jinja2.Environment()
    jinja_env.filters["percent_format"] = percent_format
    jinja_env.filters["float_format"] = float_format
    jinja_env.filters["nv_float_format"] = nv_float_format
    jinja_env.filters["date2str_format"] = date2str_format
    jinja_env.filters["datetime2str_format"] = datetime2str_format
    # 填充context数据到模板文件中
    report_tpl.render(context, jinja_env)
    #保存生成的报告文档
    report_file = os.path.join(save_dir,file_prefix+stat_date.strftime("%Y-%m-%d")+'.docx')
    report_tpl.save(report_file)